#include <gtest/gtest.h>
#include <cstdlib>
#include "FEC/Polar.h"

template <typename T, T... I>
constexpr std::array<T, sizeof...(I)> expand_sequence(std::integer_sequence<T, I...>) {
    return std::array<T, sizeof...(I)>{ I ... };
}

TEST(PolarCodeConstructionTest, CalculateDataIndices) {
    /*
    Calculated using support/MATLAB/polar/polar_consruction.m – note that they are
    one-based.
    */
    std::array<std::size_t, 512> ref_indices = {
         127,  128,  192,  224,  232,  236,  238,  239,  240,  244,  246,  247,  248,  250,  251,  252,
         253,  254,  255,  256,  320,  352,  360,  364,  366,  367,  368,  372,  374,  375,  376,  378,
         379,  380,  381,  382,  383,  384,  400,  408,  412,  414,  415,  416,  424,  426,  427,  428,
         429,  430,  431,  432,  434,  435,  436,  437,  438,  439,  440,  441,  442,  443,  444,  445,
         446,  447,  448,  452,  454,  455,  456,  458,  459,  460,  461,  462,  463,  464,  466,  467,
         468,  469,  470,  471,  472,  473,  474,  475,  476,  477,  478,  479,  480,  482,  483,  484,
         485,  486,  487,  488,  489,  490,  491,  492,  493,  494,  495,  496,  497,  498,  499,  500,
         501,  502,  503,  504,  505,  506,  507,  508,  509,  510,  511,  512,  544,  560,  568,  572,
         574,  575,  576,  592,  596,  598,  599,  600,  602,  603,  604,  605,  606,  607,  608,  612,
         614,  615,  616,  618,  619,  620,  621,  622,  623,  624,  626,  627,  628,  629,  630,  631,
         632,  633,  634,  635,  636,  637,  638,  639,  640,  648,  652,  654,  655,  656,  660,  662,
         663,  664,  666,  667,  668,  669,  670,  671,  672,  676,  678,  679,  680,  682,  683,  684,
         685,  686,  687,  688,  689,  690,  691,  692,  693,  694,  695,  696,  697,  698,  699,  700,
         701,  702,  703,  704,  708,  709,  710,  711,  712,  713,  714,  715,  716,  717,  718,  719,
         720,  721,  722,  723,  724,  725,  726,  727,  728,  729,  730,  731,  732,  733,  734,  735,
         736,  737,  738,  739,  740,  741,  742,  743,  744,  745,  746,  747,  748,  749,  750,  751,
         752,  753,  754,  755,  756,  757,  758,  759,  760,  761,  762,  763,  764,  765,  766,  767,
         768,  776,  780,  782,  783,  784,  788,  789,  790,  791,  792,  793,  794,  795,  796,  797,
         798,  799,  800,  802,  803,  804,  805,  806,  807,  808,  809,  810,  811,  812,  813,  814,
         815,  816,  817,  818,  819,  820,  821,  822,  823,  824,  825,  826,  827,  828,  829,  830,
         831,  832,  834,  835,  836,  837,  838,  839,  840,  841,  842,  843,  844,  845,  846,  847,
         848,  849,  850,  851,  852,  853,  854,  855,  856,  857,  858,  859,  860,  861,  862,  863,
         864,  865,  866,  867,  868,  869,  870,  871,  872,  873,  874,  875,  876,  877,  878,  879,
         880,  881,  882,  883,  884,  885,  886,  887,  888,  889,  890,  891,  892,  893,  894,  895,
         896,  898,  899,  900,  901,  902,  903,  904,  905,  906,  907,  908,  909,  910,  911,  912,
         913,  914,  915,  916,  917,  918,  919,  920,  921,  922,  923,  924,  925,  926,  927,  928,
         929,  930,  931,  932,  933,  934,  935,  936,  937,  938,  939,  940,  941,  942,  943,  944,
         945,  946,  947,  948,  949,  950,  951,  952,  953,  954,  955,  956,  957,  958,  959,  960,
         961,  962,  963,  964,  965,  966,  967,  968,  969,  970,  971,  972,  973,  974,  975,  976,
         977,  978,  979,  980,  981,  982,  983,  984,  985,  986,  987,  988,  989,  990,  991,  992,
         993,  994,  995,  996,  997,  998,  999, 1000, 1001, 1002, 1003, 1004, 1005, 1006, 1007, 1008,
        1009, 1010, 1011, 1012, 1013, 1014, 1015, 1016, 1017, 1018, 1019, 1020, 1021, 1022, 1023, 1024
    };

    std::array<std::size_t, 512> data_indices = expand_sequence(
        Thiemar::Polar::PolarCodeConstructor<1024, 1024, 512, -2>::data_index_sequence{});

    for (std::size_t i = 0u; i < data_indices.size(); i++) {
        EXPECT_EQ((int)ref_indices.data()[i] - 1u, (int)data_indices.data()[i]) << "Buffers differ at index " << i;
    }
}

TEST(PolarCodeConstructionTest, CalculateShortenedDataIndices768) {
    /*
    Calculated using support/MATLAB/polar/polar_consruction.m – note that they are
    one-based.
    */
    std::array<std::size_t, 512> ref_indices = {
          32,   48,   56,   60,   62,   63,   64,   80,   88,   92,   94,   95,   96,  104,  108,  110,
         111,  112,  116,  118,  119,  120,  121,  122,  123,  124,  125,  126,  127,  128,  144,  152,
         156,  158,  159,  160,  168,  172,  174,  175,  176,  180,  182,  183,  184,  185,  186,  187,
         188,  189,  190,  191,  192,  200,  204,  206,  207,  208,  212,  214,  215,  216,  217,  218,
         219,  220,  221,  222,  223,  224,  226,  227,  228,  229,  230,  231,  232,  233,  234,  235,
         236,  237,  238,  239,  240,  241,  242,  243,  244,  245,  246,  247,  248,  249,  250,  251,
         252,  253,  254,  255,  256,  272,  280,  284,  286,  287,  288,  296,  300,  302,  303,  304,
         308,  310,  311,  312,  313,  314,  315,  316,  317,  318,  319,  320,  328,  332,  334,  335,
         336,  340,  342,  343,  344,  345,  346,  347,  348,  349,  350,  351,  352,  354,  355,  356,
         357,  358,  359,  360,  361,  362,  363,  364,  365,  366,  367,  368,  369,  370,  371,  372,
         373,  374,  375,  376,  377,  378,  379,  380,  381,  382,  383,  384,  388,  390,  391,  392,
         394,  395,  396,  397,  398,  399,  400,  402,  403,  404,  405,  406,  407,  408,  409,  410,
         411,  412,  413,  414,  415,  416,  417,  418,  419,  420,  421,  422,  423,  424,  425,  426,
         427,  428,  429,  430,  431,  432,  433,  434,  435,  436,  437,  438,  439,  440,  441,  442,
         443,  444,  445,  446,  447,  448,  449,  450,  451,  452,  453,  454,  455,  456,  457,  458,
         459,  460,  461,  462,  463,  464,  465,  466,  467,  468,  469,  470,  471,  472,  473,  474,
         475,  476,  477,  478,  479,  480,  481,  482,  483,  484,  485,  486,  487,  488,  489,  490,
         491,  492,  493,  494,  495,  496,  497,  498,  499,  500,  501,  502,  503,  504,  505,  506,
         507,  508,  509,  510,  511,  512,  520,  524,  526,  527,  528,  532,  534,  536,  540,  541,
         542,  543,  544,  552,  556,  557,  558,  559,  560,  562,  563,  564,  565,  566,  567,  568,
         569,  570,  571,  572,  573,  574,  575,  576,  580,  582,  583,  584,  586,  587,  588,  589,
         590,  591,  592,  593,  594,  595,  596,  597,  598,  599,  600,  601,  602,  603,  604,  605,
         606,  607,  608,  609,  610,  611,  612,  613,  614,  615,  616,  617,  618,  619,  620,  621,
         622,  623,  624,  625,  626,  627,  628,  629,  630,  631,  632,  633,  634,  635,  636,  637,
         638,  639,  640,  644,  645,  646,  647,  648,  649,  650,  651,  652,  653,  654,  655,  656,
         657,  658,  659,  660,  661,  662,  663,  664,  665,  666,  667,  668,  669,  670,  671,  672,
         673,  674,  675,  676,  677,  678,  679,  680,  681,  682,  683,  684,  685,  686,  687,  688,
         689,  690,  691,  692,  693,  694,  695,  696,  697,  698,  699,  700,  701,  702,  703,  704,
         705,  706,  707,  708,  709,  710,  711,  712,  713,  714,  715,  716,  717,  718,  719,  720,
         721,  722,  723,  724,  725,  726,  727,  728,  729,  730,  731,  732,  733,  734,  735,  736,
         737,  738,  739,  740,  741,  742,  743,  744,  745,  746,  747,  748,  749,  750,  751,  752,
         753,  754,  755,  756,  757,  758,  759,  760,  761,  762,  763,  764,  765,  766,  767,  768
    };

    std::array<std::size_t, 512> data_indices = expand_sequence(
        Thiemar::Polar::PolarCodeConstructor<1024, 768, 512, -2>::data_index_sequence{});

    for (std::size_t i = 0u; i < data_indices.size(); i++) {
        EXPECT_EQ((int)ref_indices.data()[i] - 1u, (int)data_indices.data()[i]) << "Buffers differ at index " << i;
    }
}

TEST(PolarCodeConstructionTest, CalculateShortenedDataIndices520) {
    /*
    Calculated using support/MATLAB/polar/polar_consruction.m – note that they are
    one-based.
    */
    std::array<std::size_t, 512> ref_indices = {
           4,    6,    7,    8,   10,   11,   12,   13,   14,   15,   16,   18,   19,   20,   21,   22,
          23,   24,   25,   26,   27,   28,   29,   30,   31,   32,   33,   34,   35,   36,   37,   38,
          39,   40,   41,   42,   43,   44,   45,   46,   47,   48,   49,   50,   51,   52,   53,   54,
          55,   56,   57,   58,   59,   60,   61,   62,   63,   64,   65,   66,   67,   68,   69,   70,
          71,   72,   73,   74,   75,   76,   77,   78,   79,   80,   81,   82,   83,   84,   85,   86,
          87,   88,   89,   90,   91,   92,   93,   94,   95,   96,   97,   98,   99,  100,  101,  102,
         103,  104,  105,  106,  107,  108,  109,  110,  111,  112,  113,  114,  115,  116,  117,  118,
         119,  120,  121,  122,  123,  124,  125,  126,  127,  128,  130,  131,  132,  133,  134,  135,
         136,  137,  138,  139,  140,  141,  142,  143,  144,  145,  146,  147,  148,  149,  150,  151,
         152,  153,  154,  155,  156,  157,  158,  159,  160,  161,  162,  163,  164,  165,  166,  167,
         168,  169,  170,  171,  172,  173,  174,  175,  176,  177,  178,  179,  180,  181,  182,  183,
         184,  185,  186,  187,  188,  189,  190,  191,  192,  193,  194,  195,  196,  197,  198,  199,
         200,  201,  202,  203,  204,  205,  206,  207,  208,  209,  210,  211,  212,  213,  214,  215,
         216,  217,  218,  219,  220,  221,  222,  223,  224,  225,  226,  227,  228,  229,  230,  231,
         232,  233,  234,  235,  236,  237,  238,  239,  240,  241,  242,  243,  244,  245,  246,  247,
         248,  249,  250,  251,  252,  253,  254,  255,  256,  258,  259,  260,  261,  262,  263,  264,
         265,  266,  267,  268,  269,  270,  271,  272,  273,  274,  275,  276,  277,  278,  279,  280,
         281,  282,  283,  284,  285,  286,  287,  288,  289,  290,  291,  292,  293,  294,  295,  296,
         297,  298,  299,  300,  301,  302,  303,  304,  305,  306,  307,  308,  309,  310,  311,  312,
         313,  314,  315,  316,  317,  318,  319,  320,  321,  322,  323,  324,  325,  326,  327,  328,
         329,  330,  331,  332,  333,  334,  335,  336,  337,  338,  339,  340,  341,  342,  343,  344,
         345,  346,  347,  348,  349,  350,  351,  352,  353,  354,  355,  356,  357,  358,  359,  360,
         361,  362,  363,  364,  365,  366,  367,  368,  369,  370,  371,  372,  373,  374,  375,  376,
         377,  378,  379,  380,  381,  382,  383,  384,  385,  386,  387,  388,  389,  390,  391,  392,
         393,  394,  395,  396,  397,  398,  399,  400,  401,  402,  403,  404,  405,  406,  407,  408,
         409,  410,  411,  412,  413,  414,  415,  416,  417,  418,  419,  420,  421,  422,  423,  424,
         425,  426,  427,  428,  429,  430,  431,  432,  433,  434,  435,  436,  437,  438,  439,  440,
         441,  442,  443,  444,  445,  446,  447,  448,  449,  450,  451,  452,  453,  454,  455,  456,
         457,  458,  459,  460,  461,  462,  463,  464,  465,  466,  467,  468,  469,  470,  471,  472,
         473,  474,  475,  476,  477,  478,  479,  480,  481,  482,  483,  484,  485,  486,  487,  488,
         489,  490,  491,  492,  493,  494,  495,  496,  497,  498,  499,  500,  501,  502,  503,  504,
         505,  506,  507,  508,  509,  510,  511,  512,  513,  514,  515,  516,  517,  518,  519,  520
    };

    std::array<std::size_t, 512> data_indices = expand_sequence(
        Thiemar::Polar::PolarCodeConstructor<1024, 520, 512, -2>::data_index_sequence{});

    for (std::size_t i = 0u; i < data_indices.size(); i++) {
        EXPECT_EQ((int)ref_indices.data()[i] - 1u, (int)data_indices.data()[i]) << "Buffers differ at index " << i;
    }
}

TEST(PolarCodeConstructionTest, CalculateShortenedDataIndices512) {
    /*
    Calculated using support/MATLAB/polar/polar_consruction.m – note that they are
    one-based.
    */
    std::array<std::size_t, 512> ref_indices = {
           1,    2,    3,    4,    5,    6,    7,    8,    9,   10,   11,   12,   13,   14,   15,   16,
          17,   18,   19,   20,   21,   22,   23,   24,   25,   26,   27,   28,   29,   30,   31,   32,
          33,   34,   35,   36,   37,   38,   39,   40,   41,   42,   43,   44,   45,   46,   47,   48,
          49,   50,   51,   52,   53,   54,   55,   56,   57,   58,   59,   60,   61,   62,   63,   64,
          65,   66,   67,   68,   69,   70,   71,   72,   73,   74,   75,   76,   77,   78,   79,   80,
          81,   82,   83,   84,   85,   86,   87,   88,   89,   90,   91,   92,   93,   94,   95,   96,
          97,   98,   99,  100,  101,  102,  103,  104,  105,  106,  107,  108,  109,  110,  111,  112,
         113,  114,  115,  116,  117,  118,  119,  120,  121,  122,  123,  124,  125,  126,  127,  128,
         129,  130,  131,  132,  133,  134,  135,  136,  137,  138,  139,  140,  141,  142,  143,  144,
         145,  146,  147,  148,  149,  150,  151,  152,  153,  154,  155,  156,  157,  158,  159,  160,
         161,  162,  163,  164,  165,  166,  167,  168,  169,  170,  171,  172,  173,  174,  175,  176,
         177,  178,  179,  180,  181,  182,  183,  184,  185,  186,  187,  188,  189,  190,  191,  192,
         193,  194,  195,  196,  197,  198,  199,  200,  201,  202,  203,  204,  205,  206,  207,  208,
         209,  210,  211,  212,  213,  214,  215,  216,  217,  218,  219,  220,  221,  222,  223,  224,
         225,  226,  227,  228,  229,  230,  231,  232,  233,  234,  235,  236,  237,  238,  239,  240,
         241,  242,  243,  244,  245,  246,  247,  248,  249,  250,  251,  252,  253,  254,  255,  256,
         257,  258,  259,  260,  261,  262,  263,  264,  265,  266,  267,  268,  269,  270,  271,  272,
         273,  274,  275,  276,  277,  278,  279,  280,  281,  282,  283,  284,  285,  286,  287,  288,
         289,  290,  291,  292,  293,  294,  295,  296,  297,  298,  299,  300,  301,  302,  303,  304,
         305,  306,  307,  308,  309,  310,  311,  312,  313,  314,  315,  316,  317,  318,  319,  320,
         321,  322,  323,  324,  325,  326,  327,  328,  329,  330,  331,  332,  333,  334,  335,  336,
         337,  338,  339,  340,  341,  342,  343,  344,  345,  346,  347,  348,  349,  350,  351,  352,
         353,  354,  355,  356,  357,  358,  359,  360,  361,  362,  363,  364,  365,  366,  367,  368,
         369,  370,  371,  372,  373,  374,  375,  376,  377,  378,  379,  380,  381,  382,  383,  384,
         385,  386,  387,  388,  389,  390,  391,  392,  393,  394,  395,  396,  397,  398,  399,  400,
         401,  402,  403,  404,  405,  406,  407,  408,  409,  410,  411,  412,  413,  414,  415,  416,
         417,  418,  419,  420,  421,  422,  423,  424,  425,  426,  427,  428,  429,  430,  431,  432,
         433,  434,  435,  436,  437,  438,  439,  440,  441,  442,  443,  444,  445,  446,  447,  448,
         449,  450,  451,  452,  453,  454,  455,  456,  457,  458,  459,  460,  461,  462,  463,  464,
         465,  466,  467,  468,  469,  470,  471,  472,  473,  474,  475,  476,  477,  478,  479,  480,
         481,  482,  483,  484,  485,  486,  487,  488,  489,  490,  491,  492,  493,  494,  495,  496,
         497,  498,  499,  500,  501,  502,  503,  504,  505,  506,  507,  508,  509,  510,  511,  512
    };

    std::array<std::size_t, 512> data_indices = expand_sequence(
        Thiemar::Polar::PolarCodeConstructor<1024, 512, 512, -2>::data_index_sequence{});

    for (std::size_t i = 0u; i < data_indices.size(); i++) {
        EXPECT_EQ((int)ref_indices.data()[i] - 1u, (int)data_indices.data()[i]) << "Buffers differ at index " << i;
    }
}
